import numpy as np
from LTE import constant as c
import dsp.dsp as DSP

class Phy:
    @staticmethod
    def ofdm_tx(sf_symbols,nof_prb):
        symbol_sz = c.LTE_SYMBOL_SZ[nof_prb]
        cp_len0   = c.LTE_CP_LEN_NORM(0, symbol_sz)
        cp_len1   = c.LTE_CP_LEN_NORM(1, symbol_sz)
        slot_len  = c.LTE_SLOT_LEN(symbol_sz)
        sequences = np.zeros((c.LTE_NOF_SF_X_FRAME,slot_len),np.complex64)
        n = nof_prb * c.LTE_NRE
        for sf in range(c.LTE_NOF_SF_X_FRAME):
            for symbidx in range(c.LTE_CP_NORM_SF_NSYMB):
                offset  = cp_len0 + (symbol_sz + cp_len1) * symbidx

                if symbidx >= c.LTE_CP_NORM_NSYMB:
                    offset += cp_len0 - cp_len1

                tmp = np.zeros((symbol_sz,),np.complex64)
                fsym = sf_symbols[sf,symbidx,:]
                hre = n//2
                tmp[1:hre+1] = fsym[hre:]
                tmp[-hre:] = fsym[:hre]
                tsym = np.fft.ifft(tmp)
                sequences[sf,offset:offset+symbol_sz] = tsym[:]
                cp = cp_len0 if symbidx%c.LTE_CP_NORM_NSYMB == 0 else cp_len1
                sequences[sf,offset-cp:offset] = tsym[-cp:]
        sequences *= symbol_sz/(4*2*nof_prb)    
        return sequences

    @staticmethod
    def ofdm_rx(rfd,nof_prb):
        symbol_sz = c.LTE_SYMBOL_SZ[nof_prb]
        sf_symbols = np.zeros((c.LTE_NOF_SF_X_FRAME,c.LTE_CP_NORM_SF_NSYMB,symbol_sz),dtype=complex)
        cp_len0   = c.LTE_CP_LEN_NORM(0, symbol_sz)
        cp_len1   = c.LTE_CP_LEN_NORM(1, symbol_sz)
        slot_len  = c.LTE_SLOT_LEN(symbol_sz)
        sequences = np.reshape(rfd,(c.LTE_NOF_SF_X_FRAME,slot_len))
        for sf in range(c.LTE_NOF_SF_X_FRAME):
            for symbidx in range(c.LTE_CP_NORM_SF_NSYMB):
                offset  = cp_len0 + (symbol_sz + cp_len1) * symbidx

                if symbidx >= c.LTE_CP_NORM_NSYMB:
                    offset += cp_len0 - cp_len1

                tmp = sequences[sf,offset:offset+symbol_sz]
                tmp[-3:] = sequences[sf,offset-3:offset]
                sf_symbols[sf,symbidx,:]=np.fft.fft(tmp)
        return sf_symbols
    
    @staticmethod
    def cs_slot_measure(rfd,ref,r):
        res = np.zeros((c.LTE_NOF_SF_X_FRAME,len(r)),dtype=complex)
        for i in range(c.LTE_NOF_SF_X_FRAME):
            res[i,:] = DSP.cycle_corr(rfd,ref[i,:],r)
        return res

    @staticmethod
    def cs_frame_measure(rfd,ref,r):
        sm = Phy.cs_slot_measure(rfd,ref,r)
        return np.sum(sm,axis=1)
    
    @staticmethod
    def pick_pilot(rx_symbols,cs_tab):
        res = np.zeros(128,dtype=complex)
        for (sf,sym,idx) in cs_tab[0]:
            res[idx] += rx_symbols[sf,sym,idx]
        for (sf,sym,idx) in cs_tab[1]:
            res[idx] -= rx_symbols[sf,sym,idx]
        for (sf,sym,idx) in cs_tab[2]:
            res[idx] -= 1j*rx_symbols[sf,sym,idx]
        for (sf,sym,idx) in cs_tab[3]:
            res[idx] += 1j*rx_symbols[sf,sym,idx]
        return res*(1+1j)*np.sqrt(0.5)